package org.btik.light.server.dev.sensitive;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.btik.light.meta.exception.LightException;
import org.btik.light.server.platform.common.api.bean.sensitive.DhPublicKey;
import org.btik.light.server.platform.common.api.service.SensitiveDataService;
import org.btik.light.server.platform.common.timeout.TimeOutCache;
import org.btik.light.tool.DhUtil;
import org.btik.light.tool.UuidUtil;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import javax.crypto.interfaces.DHPublicKey;
import javax.crypto.spec.DHParameterSpec;
import java.math.BigInteger;
import java.security.NoSuchAlgorithmException;
import java.util.Base64;

import static org.btik.light.tool.ByteUtil.encode2Base64Str;

/**
 * @author lustre
 * @since 2023/6/7 10:53
 */
@Service
public class SensitiveDataServiceImpl implements SensitiveDataService {

    private final Log log = LogFactory.getLog(SensitiveDataServiceImpl.class);

    private static final String SESSION_NAMESPACE = "DhKeyPairSession";

    @Value("${sensitive.dh-keypair-timeout:300000}")
    private long dhKeyPairSessionTimeOut;

    @Autowired
    private TimeOutCache timeOutCache;

    @Override
    public DhPublicKey exchangeRequest() {
        try {
            DhUtil.DhKeyPair dhKeyPair = DhUtil.initKey();
            DHPublicKey publicKey = dhKeyPair.getPublicKey();
            DHParameterSpec params = publicKey.getParams();

            String exchangeSessionId = UuidUtil.uuidString(8);
            timeOutCache.save(SESSION_NAMESPACE,exchangeSessionId, dhKeyPair, dhKeyPairSessionTimeOut);
            return new DhPublicKey(
                    encode2Base64Str(publicKey.getY().toByteArray()),
                    encode2Base64Str(params.getP().toByteArray()),
                    encode2Base64Str(params.getG().toByteArray()),
                    exchangeSessionId
            );
        } catch (NoSuchAlgorithmException e) {
            throw new LightException(e);
        }

    }

    @Override
    public DhUtil.DhKeyPair generateKeyPair() {
        try {
            return DhUtil.initKey();
        } catch (NoSuchAlgorithmException e) {
            throw new LightException(e);
        }
    }

    @Override
    public byte[] exchangeFinal(String publicKey, String exchangeSessionId) {
        try {
            DhUtil.DhKeyPair dhKeyPair = timeOutCache.getObj(SESSION_NAMESPACE, exchangeSessionId);
            return exchangeFinal(publicKey, dhKeyPair);
        } finally {
            timeOutCache.removeObj(SESSION_NAMESPACE, exchangeSessionId);
        }
    }

    @Override
    public byte[] exchangeFinal(String publicKey, DhUtil.DhKeyPair dhKeyPair) {
        if (dhKeyPair == null) {
            return null;
        }
        byte[] pkArray = Base64.getDecoder().decode(publicKey);
        // 在sun.security.util.KeyUtil.validateDHPublicKey(java.math.BigInteger, java.math.BigInteger, java.math.BigInteger)
        // 该方法会检查是否为正数 ，由于公钥含可能第一位的byte可能超过127，不指定正负情况下
        // 符号位会被用作符号 会导致检查报错 Diffie-Hellman public key is too small
        BigInteger k = new BigInteger(1, pkArray);
        DHParameterSpec params = dhKeyPair.getPublicKey().getParams();
        byte[] secretKey;
        try {
            secretKey = DhUtil.keyAgree(k, params.getP(), params.getG(), dhKeyPair.getPrivateKey());
        } catch (Exception e) {
            log.error("", e);
            return null;
        }
        return secretKey;
    }
}
